#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue

template<typename typeT> class KernelReshape {
public:
    __aicore__ inline KernelReshape() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR shape, GM_ADDR y, 
                                uint32_t totallength, 
                                uint32_t tilelength) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->totallength = totallength;
        this->tilelength = tilelength;
        this->alignnum = 32 / sizeof(typeT);

        xGm.SetGlobalBuffer((__gm__ typeT*)x, this->totallength + this->alignnum);
        yGm.SetGlobalBuffer((__gm__ typeT*)y, this->totallength + this->alignnum);

        // pipe alloc memory to queue, the unit is Bytes
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tilelength * sizeof(typeT));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tilelength * sizeof(typeT));
        // pipe.InitBuffer(tmp1, this->tilelength * sizeof(float));
        // pipe.InitBuffer(tmp2, this->tilelength * sizeof(float));
        // pipe.InitBuffer(tmp3, this->tilelength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        bool istail = false;
        int32_t loopCount = this->totallength / this->tilelength;
        if(loopCount == 0) {
            istail = true;
            uint32_t lengthalign = (this->totallength + this->alignnum - 1) / this->alignnum * this->alignnum;
            CopyIn(0, lengthalign, istail);
            CopyOut(0, lengthalign, istail);
        } else {
            for (int32_t i = 0; i < loopCount; i++) {
                istail = false;
                CopyIn(i, this->tilelength, istail);
                CopyOut(i, this->tilelength, istail);
            } 
            if (this->totallength % this->tilelength > 0) {
                istail = true;
                uint32_t length = this->totallength % this->tilelength;
                uint32_t lengthalign = (length + this->alignnum - 1) / this->alignnum * this->alignnum;
                CopyIn(loopCount, lengthalign, istail);
                CopyOut(loopCount, lengthalign, istail);
            }   
        }        
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length, bool istail) {
        LocalTensor<typeT> xLocal = inQueueX.AllocTensor<typeT>();
        DataCopy(xLocal, xGm[progress * this->tilelength], length);
        // if(istail == false) { 
        //     DataCopy(xLocal, xGm[progress * this->tilelength], length);
        // } else {
        //     int32_t startpos = progress * this->tilelength;
        //     uint32_t blockLen = length * sizeof(DTYPE_X);
        //     DataCopyExtParams copyParams{1, blockLen, 0, 0, 0}; 
        //     DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        //     DataCopyPad(xLocal, xGm[startpos], copyParams, padParams);
        // }    
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length, bool istail) {
        LocalTensor<typeT> xLocal = inQueueX.DeQue<typeT>();
        DataCopy(yGm[progress * this->tilelength], xLocal, length);
        // if(istail == false) { 
        //     DataCopy(yGm[progress * this->tilelength], xLocal, length);
        // } else {
        //     int32_t startpos = progress * this->tilelength;
        //     uint32_t blockLen = length * sizeof(DTYPE_X);
        //     DataCopyExtParams copyParams{1, blockLen, 0, 0, 0};
        //     DataCopyPad(yGm[startpos], xLocal, copyParams); 
        // }    
        inQueueX.FreeTensor(xLocal);
    }

private:
    TPipe pipe;
    //TBuf<QuePosition::VECCALC> tmp1, tmp2, tmp3;
    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> inQueueX;
    //TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<typeT> xGm;
    GlobalTensor<typeT> yGm;
    uint32_t totallength;
    uint32_t tilelength;
    uint32_t alignnum;

};

extern "C" __global__ __aicore__ void reshape(GM_ADDR x, GM_ADDR shape, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelReshape<DTYPE_X> op;
    op.Init(x, shape, y, tiling_data.totallength, tiling_data.tilelength);
    op.Process();
}